#!/usr/bin/env python3
# -*- coding: utf-8 -*

import os
import sys
import yaml
import argparse

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from package.utility import ensure_exists

from create_sub_file.create_subfile_main import *
from get_free_nodes.get_free_nodes import get_free_nodes_injs

# TODO:
# 2. 某些参数可以为空，方便能形成简单的测试文件
# 3. 将代码梳理一遍

# set parser
def make_parser(): 
    parser = argparse.ArgumentParser(
        prog="qsub to pbs schedule",
        description="create pbs script and sub to pbs.")
    parser.add_argument(
            "CONFIGFILE",
            metavar="CONFIGFILE",
            default="sugon.yaml",
            help="Config file of setting info"
    )
    return parser

def set_work_dir(myconfig):
    workdir = myconfig["workdir"]
    username = os.popen("echo $USER").read().strip("\n")
    machinename = myconfig["subinfo"]["machine_name"]
    workdir = os.path.join(workdir, f"{machinename}_{username}")
    os.system("mkdir -p {0}".format(workdir))
    os.chdir(workdir)

def read_myconfig(configpath):
    relconfiggpath = os.path.abspath(configpath)
    ensure_exists(relconfiggpath)
    with open(relconfiggpath, "r") as f:
        myconfig = yaml.safe_load(f)
    return myconfig

def checkconfig(myconfig):
    if myconfig["workdir"] is None:
        print("Workdir empty!")
        exit(1)
    if myconfig["program"] is None:
        print("Workdir empty!")
        exit(1)
    if myconfig["amg_parameter"]:
        check_amg_parameter(myconfig["amg_parameter"])
    if myconfig["subinfo"]:
        check_subinfo(myconfig["subinfo"])

def check_amg_parameter(amg_parameter):
    numcase = len(amg_parameter["problem"]["jxpamg_matrix_dir"])
    numrhs  = len(amg_parameter["problem"]["rhs_dir"])
    if numcase != numrhs:
        print("Number of matrix is not equal to numver of rhs!")
        exit(1)

def check_subinfo(subfinfo):
    np_start = subfinfo["processinfo"]["np_start"]
    np_end   =  subfinfo["processinfo"]["np_end"]
    if np_start > np_end or np_start < 0 or np_end < 0:
        print("Processinfo erro")
        exit(1)

def main():
    parser = make_parser()
    args = parser.parse_args()
    myconfig = read_myconfig(args.CONFIGFILE)
    #TODO:checkconfig需要完善
    checkconfig(myconfig)
    set_work_dir(myconfig)
    get_free_nodes_injs(myconfig)
    create_sub_file(myconfig)


if __name__=="__main__":
    main()
